# 图片操作
from PIL import Image
import numpy as np
import cv2
from zz import zz_return, zz_str


def find_img(img_src, img_target, threshold=0.9, get_index=0) -> zz_return.Return:
    """
    目标图在原图的位置,返回x,y
    :param img_src: 原图
    :param img_target: 目标图
    :param threshold: 设置阈值
    :param get_index: 匹配成功的数组下标，默认获取第1个
    :return:
    """
    img_arr = load_img(img_src, cv2.IMREAD_GRAYSCALE)
    img_target_arr = load_img(img_target, cv2.IMREAD_GRAYSCALE)  # 以灰度模式读取
    # print(img_target_arr)
    # 获取模板的宽度和高度
    w, h = img_target_arr.shape[::-1]

    # 归一化相关系数匹配法
    res = cv2.matchTemplate(img_arr, img_target_arr, cv2.TM_CCOEFF_NORMED)  # 执行模板匹配
    loc = np.where(res >= threshold)  # 查找最佳匹配位置

    # 使用平方差匹配算法
    # res = cv2.matchTemplate(img_arr, img_target_arr, cv2.TM_SQDIFF)
    # loc = np.where(res < threshold)
    loc = list(zip(*loc[::-1]))  # 将两个1维组数，合并变成（x,y）
    loc = find_img_location_remove(loc)
    # print(loc)
    if len(loc) <= get_index:
        return zz_return.of(20029, "未匹配到位置")

    loc_list = []  # 所有坐标
    loc_list_center = []  # 所有中心标签
    for item in loc:
        x, y = item
        x = int(x)
        y = int(y)
        center_x = x + int(w / 2)
        center_y = y + int(h / 2)
        loc_list.append([x, y])
        loc_list_center.append([center_x, center_y])

    r = zz_return.of(0)
    r.set_data("first", loc_list[0])
    r.set_data("center", loc_list_center[0])
    r.set_data("center_list", loc_list_center)

    # 返回匹配位置
    return r


def find_img_rectangle(img_src, img_target, threshold=0.9):
    """
    目标图在原图的位置,并画出红框
    :param img_src: 原图
    :param img_target: 目标图
    :param threshold: 设置阈值
    """
    # 图片灰度模式读取,为了速度更快
    img_arr = load_img(img_src, cv2.IMREAD_GRAYSCALE)
    img_target_arr = cv2.imread(img_target, cv2.IMREAD_GRAYSCALE)  # 以灰度模式读取
    # print(img_target_arr)
    # 获取模板的宽度和高度
    w, h = img_target_arr.shape[::-1]

    # 归一化相关系数匹配法
    res = cv2.matchTemplate(img_arr, img_target_arr, cv2.TM_CCOEFF_NORMED)  # 执行模板匹配
    # print(res)
    loc = np.where(res >= threshold)  # 查找最佳匹配位置

    # 使用平方差匹配算法
    # res = cv2.matchTemplate(img_arr, img_target_arr, cv2.TM_SQDIFF)
    # loc = np.where(res < threshold)
    loc = list(zip(*loc[::-1]))  # 将两个1维组数，合并变成（x,y）
    loc = find_img_location_remove(loc)
    # print(loc)
    if len(loc) <= 0:
        raise ValueError("未匹配到位置")

    img_save = np.array(cv2.imread(img_src))
    for item in loc:
        print(item)
        x, y = item
        # 画框
        cv2.rectangle(img_save, (x, y), (x + w, y + h), (0, 0, 255), 1)

    return img_save


def find_img_rectangle_show(img_src, img_target, threshold=0.9):
    img_save = find_img_rectangle(img_src, img_target, threshold)

    cv2.namedWindow("preview", cv2.WINDOW_NORMAL)

    # src_w, src_h = img_arr.shape[::-1]
    # cv2.resizeWindow("preview", src_w, src_h)
    cv2.imshow("preview", img_save)
    cv2.waitKey(0)
    cv2.destroyAllWindows()


def find_img_location_remove(arr, threshold=10):
    """
    去掉一些误差小的坐标
    """
    if len(arr) <= 0:
        return arr

    arr_new = []
    arr_new_index = (0, 0)
    for item in arr:
        add = False
        if abs(item[0] - arr_new_index[0]) > threshold or abs(item[1] - arr_new_index[1]) > threshold:
            add = True
            arr_new_index = item
        if add:
            arr_new.append(item)
    return arr_new


def color_match(img_src, img_cat_top, img_cat_bottom, colors: list, threshold=10):
    """
    匹配颜块位置
    :param img_src: 图片路径
    :param img_cat_top: 坐标xy，例：(100,455)
    :param img_cat_bottom: 底部坐标xy，存在此值，
    :param colors 颜色数组
    :param threshold 误差
    """
    image = load_img(img_src)
    image = image[:, :, [2, 1, 0]]
    # print(image)

    # 获取柜阵的左下角和右上角坐标
    left, top = min(img_cat_top[0], img_cat_bottom[0]), min(img_cat_top[1], img_cat_bottom[1])
    right, bottom = max(img_cat_top[0], img_cat_bottom[0]), max(img_cat_top[1], img_cat_bottom[1])

    # 获取指定坐标，柜阵
    img_cat_arr = []
    for y in range(top, bottom + 1):
        row = []
        for x in range(left, right + 1):
            rgb = image[y][x]
            # print(f"x,y = {x},{y}; rgb = {rgb}")
            row.append(rgb)
        img_cat_arr.append(row)

    img_cat_arr = np.array(img_cat_arr)
    colors_to = color_threshold(colors, threshold)
    color_arr = np.array(colors_to)
    # print(img_cat_arr)

    # positions = np.where((img_cat_arr is color_arr).all(axis=2))
    condition = np.isin(img_cat_arr, color_arr).all(axis=2)
    # print(condition)
    row_x, row_y = np.where(condition)
    # print(row_x, row_y)
    positions = list(zip(row_x, row_y))
    if len(positions) <= 0:
        return zz_return.of(20164, "未匹配到坐标")
    xy_list = []
    # print(positions)
    for y, x in positions:
        xx = int(x + left)
        yy = int(y + top)
        # print(xx, yy)
        xy_list.append((xx, yy))
        # print(f"颜色为{image[yy][xx]}的像素位置为({xx}, {yy})")
    r = zz_return.of(0)
    r.set_data("list", xy_list)
    return r


def color_threshold(color: list, threshold):
    """
    颜色列表设置误差
    """
    new_color = []
    for rgb in color:
        r, g, b = rgb
        new_color.append(rgb)
        [new_color.append((r - i, g, b)) for i in range(1, threshold + 1)]
        [new_color.append((r + i, g, b)) for i in range(1, threshold + 1)]

        [new_color.append((r, g - i, b)) for i in range(1, threshold + 1)]
        [new_color.append((r, g + i, b)) for i in range(1, threshold + 1)]

        [new_color.append((r, g, b - i)) for i in range(1, threshold + 1)]
        [new_color.append((r, g, b + i)) for i in range(1, threshold + 1)]
    return new_color


def load_img(img_src, flags=cv2.IMREAD_UNCHANGED):
    """
    图片加载
    :param img_src:
    :flags img_src: 原图cv2.IMREAD_UNCHANGED，灰度图片cv2.IMREAD_GRAYSCALE
    """
    if isinstance(img_src, np.ndarray):
        image = img_src
    else:
        image = cv2.imread(img_src, flags)
    return image
